#include <xy/render/image.h>

#define STB_IMAGE_IMPLEMENTATION
#include <stb/stb_image.h>

#define STB_IMAGE_WRITE_IMPLEMENTATION
#include <stb/stb_image_write.h>

namespace xy
{

namespace render
{

color_t::color_t(std::uint32_t argb)
    : bgra{(argb >> 0) & 0xff, (argb >> 8) & 0xff, (argb >> 16) & 0xff, (argb >> 24) & 0xff}
{
}

color_t::color_t(std::uint8_t R, std::uint8_t G, std::uint8_t B, std::uint8_t A) : bgra{B, G, R, A}
{
}

color_t::color_t(const std::uint8_t *p, std::size_t bpp)
{
    std::fill(bgra, bgra + 4, 255);
    std::memcpy(bgra, p, bpp);
}

color_t color_t::operator+(const color_t &other) const
{
    // alpha 取后来者
    return color_t(std::min(bgra[2] + other.bgra[2], 255), std::min(bgra[1] + other.bgra[1], 255),
                   std::min(bgra[0] + other.bgra[0], 255), other.bgra[3]);
}

color_t color_t::operator*(float t) const
{
    // alpha 值不参与运算
    return color_t(bgra[2] * t, bgra[1] * t, bgra[0] * t, bgra[3]);
}

color_t::operator std::uint32_t() const
{
    return (bgra[3] << 24) | (bgra[2] << 16) | (bgra[1] << 8) | bgra[0];
}

color_t color_t::alpha(const color_t &fg, const color_t &bg)
{
    float fgAlpha = fg.bgra[3] / 255.0f;
    float bgAlpha = bg.bgra[3] / 255.0f;
    float dstAlpha = 1.0f - (1.0f - fgAlpha) * (1.0f - bgAlpha);
    return color_t(bg.bgra[2] * (1.0f - fgAlpha) + fg.bgra[2] * fgAlpha,
                   bg.bgra[1] * (1.0f - fgAlpha) + fg.bgra[1] * fgAlpha,
                   bg.bgra[0] * (1.0f - fgAlpha) + fg.bgra[0] * fgAlpha, 255 * dstAlpha);
}

image_t::image_t() : m_impl(std::make_shared<_impl>())
{
}

image_t::image_t(std::size_t width, std::size_t height, std::size_t bpp)
    : m_impl(std::make_shared<_impl>(width, height, bpp))
{
}

image_t::_impl *image_t::operator->() const noexcept
{
    return m_impl.get();
}

color_t image_t::get(std::size_t x, std::size_t y) const
{
#ifdef _DEBUG
    if (x < 0 || x >= m_impl->width || y < 0 || y >= m_impl->height)
        throw std::out_of_range("pixel out of range");
#endif
    return color_t(m_impl->data.data() + (x + y * m_impl->width) * m_impl->bpp, m_impl->bpp);
}

void image_t::set(std::size_t x, std::size_t y, const color_t &c)
{
#ifdef _DEBUG
    if (x < 0 || x >= m_impl->width || y < 0 || y >= m_impl->height)
        throw std::out_of_range("pixel out of range");
#endif
    std::memcpy(m_impl->data.data() + (x + y * m_impl->width) * m_impl->bpp, c.bgra, m_impl->bpp);
}

color_t image_t::safe_get(std::size_t x, std::size_t y, const color_t &c) const
{
    if (x < 0 || x >= m_impl->width || y < 0 || y >= m_impl->height)
        return c;
    return color_t(m_impl->data.data() + (x + y * m_impl->width) * m_impl->bpp, m_impl->bpp);
}

void image_t::safe_set(std::size_t x, std::size_t y, const color_t &c)
{
    if (x < 0 || x >= m_impl->width || y < 0 || y >= m_impl->height)
        return;
    std::memcpy(m_impl->data.data() + (x + y * m_impl->width) * m_impl->bpp, c.bgra, m_impl->bpp);
}

void image_t::invert_rgb()
{
    // 使用指针防止索引溢出问题
    std::uint8_t *pixel = m_impl->data.data();
    for (std::size_t i = 0; i < m_impl->width; i++)
    {
        for (std::size_t j = 0; j < m_impl->height; j++)
        {
            std::swap(*(pixel + 0), *(pixel + 2));
            pixel += m_impl->bpp;
        }
    }
}

void image_t::flip_horizontally()
{
    std::size_t half = m_impl->width >> 1;
    for (std::size_t i = 0; i < half; i++)
        for (std::size_t j = 0; j < m_impl->height; j++)
            for (std::size_t b = 0; b < m_impl->bpp; b++)
                std::swap(m_impl->data[(i + j * m_impl->width) * m_impl->bpp + b],
                          m_impl->data[(m_impl->width - 1 - i + j * m_impl->width) * m_impl->bpp + b]);
}

void image_t::flip_vertically()
{
    std::size_t half = m_impl->height >> 1;
    for (std::size_t i = 0; i < m_impl->width; i++)
        for (std::size_t j = 0; j < half; j++)
            for (std::size_t b = 0; b < m_impl->bpp; b++)
                std::swap(m_impl->data[(i + j * m_impl->width) * m_impl->bpp + b],
                          m_impl->data[(i + (m_impl->height - 1 - j) * m_impl->width) * m_impl->bpp + b]);
}

void image_t::resize(std::size_t width, std::size_t height)
{
    m_impl->width = width;
    m_impl->height = height;
    m_impl->data.resize(width * height * m_impl->bpp, 0);
}

void image_t::read(std::string filename)
{
    // 加载图片前翻转图片
    stbi_set_flip_vertically_on_load(true);

    int width, height, nrComponents;
    unsigned char *data = stbi_load(filename.c_str(), &width, &height, &nrComponents, 0);
    if (data)
    {
        m_impl->width = width;
        m_impl->height = height;
        m_impl->bpp = nrComponents;
        m_impl->data.resize(width * height * nrComponents, 0);

        std::memcpy(m_impl->data.data(), data, width * height * nrComponents);
        stbi_image_free(data);

        // 注意翻转 rgb
        invert_rgb();
    }
    else
    {
        std::cerr << "Image failed to load at path: " << filename;
        stbi_image_free(data);
    }
}

void image_t::write(std::string filename) const
{
    stbi_flip_vertically_on_write(true);

    const_cast<image_t *>(this)->invert_rgb();
    int success_png = stbi_write_png(filename.c_str(), m_impl->width, m_impl->height, m_impl->bpp, m_impl->data.data(),
                                     m_impl->width * m_impl->bpp);
    const_cast<image_t *>(this)->invert_rgb();

    if (!success_png)
        std::cerr << "Image failed to save at path: " << filename;
}

image_t image_t::half_image() const
{
    static constexpr float filter[3][3] = {
        {1.0f / 16, 2.0f / 16, 1.0f / 16}, {2.0f / 16, 4.0f / 16, 2.0f / 16}, {1.0f / 16, 2.0f / 16, 1.0f / 16}};

    image_t result(m_impl->width >> 1, m_impl->height >> 1, m_impl->bpp);
    for (std::size_t i = 0; i < result.m_impl->width; i++)
    {
        for (std::size_t j = 0; j < result.m_impl->height; j++)
        {
            color_t color(0, 0, 0, 255);
            for (std::size_t k = 0; k < 3; k++)
            {
                for (std::size_t l = 0; l < 3; l++)
                {
                    std::size_t x = i * 2 + k - 1;
                    std::size_t y = j * 2 + l - 1;
                    if (x >= 0 && x < m_impl->width && y >= 0 && y < m_impl->height)
                        color = color + get(x, y) * filter[k][l];
                }
            }
            result.set(i, j, color);
        }
    }
    return result;
}

} // namespace render

} // namespace xy
